# encoding=utf-8
import torch
from transformers import BertTokenizer, BertForSequenceClassification

# 加载预训练的模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

# 定义要分类的句子
text = "This is a great movie. I really enjoyed it!"

# 对句子进行分词，并加上特殊标记
tokens = tokenizer.encode_plus(text, max_length=128, padding='max_length', truncation=True, return_tensors='pt')

# 使用模型进行预测
outputs = model(tokens['input_ids'], tokens['attention_mask'])
predicted_class = torch.argmax(outputs[0]).item()

# 输出预测结果
if predicted_class == 1:
    print("Positive sentiment")
else:
    print("Negative sentiment")
